{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lucid: Cluster Predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import random\n",
    "\n",
    "from sklearn import preprocessing, metrics\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.neural_network import MLPRegressor\n",
    "import lightgbm as lgb\n",
    "import xgboost as xgb\n",
    "from primo.model import PrimoRegressor\n",
    "\n",
    "sns.set_style(\"ticks\")\n",
    "font = {\n",
    "    \"font.family\": \"Roboto\",\n",
    "    \"font.size\": 12,\n",
    "}\n",
    "sns.set_style(font)\n",
    "paper_rc = {\n",
    "    \"lines.linewidth\": 3,\n",
    "    \"lines.markersize\": 10,\n",
    "}\n",
    "sns.set_context(\"paper\", font_scale=1.8, rc=paper_rc)\n",
    "current_palette = sns.color_palette()\n",
    "\n",
    "pd.set_option(\"display.max_columns\", None)\n",
    "\n",
    "def set_seed(seed):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "\n",
    "\n",
    "seed = 123\n",
    "set_seed(seed)\n",
    "\n",
    "cluster_list = [\"Venus\", \"Saturn\", \"Philly\"]\n",
    "cluster = cluster_list[0]\n",
    "df = pd.read_csv(\n",
    "    f\"../data/{cluster}/cluster_throughput.csv\",\n",
    "    parse_dates=[\"time\"],\n",
    "    index_col=\"time\",\n",
    "    usecols=[\"time\", \"submit_gpu_job\", \"submit_gpu_num\"],\n",
    ")\n",
    "result = pd.DataFrame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def soft_avg(data, t):\n",
    "    avg = 0.5 * data.transform(lambda x: x.shift(t))\n",
    "    avg += 0.25 * (data.transform(lambda x: x.shift(t + 1)) + data.transform(lambda x: x.shift(t - 1)))\n",
    "    return avg\n",
    "\n",
    "\n",
    "def rolling_data(df, item):\n",
    "    if item == \"submit_gpu_job\":\n",
    "        short = \"njob\"\n",
    "    elif item == \"submit_gpu_num\":\n",
    "        short = \"ngpu\"\n",
    "    else:\n",
    "        raise ValueError\n",
    "\n",
    "    df[f\"shift_1h_{short}\"] = df[item].transform(lambda x: x.shift(6))  # 1 Hour Ago\n",
    "    df[f\"shift_3h_{short}\"] = df[item].transform(lambda x: x.shift(18))  # 3 Hour Ago\n",
    "    df[f\"shift_1d_{short}\"] = df[item].transform(lambda x: x.shift(144))  # 1 Day Ago\n",
    "\n",
    "    df[f\"soft_1h_{short}\"] = soft_avg(df[item], 6)  # Soft 1 Hour Ago\n",
    "    df[f\"soft_3h_{short}\"] = soft_avg(df[item], 18)  # Soft 3 Hour Ago\n",
    "    df[f\"soft_1d_{short}\"] = soft_avg(df[item], 144)  # Soft 1 Day Ago\n",
    "\n",
    "    df[f\"roll_mean_1h_{short}\"] = df[item].transform(lambda x: x.shift(6).rolling(6).mean())  # 1 Hour Ago\n",
    "    df[f\"roll_median_1h_{short}\"] = df[item].transform(lambda x: x.shift(6).rolling(6).median())  # 1 Hour Ago\n",
    "    df[f\"roll_std_1h_{short}\"] = df[item].transform(lambda x: x.shift(6).rolling(6).std())  # 1 Hour Ago\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "def feature_engineering(df):\n",
    "    data = df.reset_index()\n",
    "\n",
    "    \"\"\"Time Features\"\"\"\n",
    "    time_features = [\"month\", \"day\", \"hour\", \"minute\", \"dayofweek\", \"dayofyear\"]\n",
    "    for tf in time_features:\n",
    "        data[tf] = getattr(data[\"time\"].dt, tf).astype(np.int16)\n",
    "    data[\"week\"] = data[\"time\"].dt.isocalendar().week.astype(np.int16)  # weekofyear\n",
    "\n",
    "    \"\"\"Job Num & GPU Num Rolling\"\"\"\n",
    "    data = rolling_data(data, \"submit_gpu_job\")\n",
    "    data = rolling_data(data, \"submit_gpu_num\")\n",
    "\n",
    "    return data\n",
    "\n",
    "\n",
    "def plot_predict(pred, test, item, cluster, ymax=100, save=False):\n",
    "    fig, ax = plt.subplots(figsize=(8, 3))\n",
    "    x = np.arange(len(test))\n",
    "\n",
    "    if \"job\" in item:\n",
    "        str_y = \"Job Submission\"  # per 10 minutes\"\n",
    "    elif \"num\" in item:\n",
    "        str_y = \"GPU Submission\"  # per 10 minutes\"\n",
    "    else:\n",
    "        raise ValueError\n",
    "\n",
    "    if cluster == \"Philly\":\n",
    "        str_month = \"November\"\n",
    "    else:\n",
    "        str_month = \"September\"\n",
    "\n",
    "    ax.plot(x, test[item].values, linestyle=\"--\", alpha=0.9, label=\"Real\", linewidth=2)\n",
    "    ax.plot(x, pred, linestyle=\"-\", alpha=0.9, label=\"Prediction\", linewidth=2)\n",
    "    ax.set_xlabel(f\"Date in {str_month} ({cluster})\")\n",
    "    ax.set_ylabel(str_y)\n",
    "\n",
    "    tick_interval = 2 * 144  # Two days\n",
    "    ax.set_xticks(x[::tick_interval])\n",
    "    ax.set_xticklabels(test[\"time\"].dt.day[::tick_interval])\n",
    "    ax.set_xlim(0, len(test) + 1)\n",
    "    ax.set_ylim(0, ymax)\n",
    "    ax.grid(axis=\"y\", linestyle=\":\")\n",
    "    ax.legend()\n",
    "    sns.despine()\n",
    "\n",
    "    if save:\n",
    "        fig.savefig(f\"./{cluster}_throughput_{str_y}.pdf\", bbox_inches=\"tight\")\n",
    "\n",
    "\n",
    "data = feature_engineering(df)\n",
    "\n",
    "if cluster == \"Philly\":\n",
    "    month = 11\n",
    "else:\n",
    "    month = 9\n",
    "\n",
    "train = data[data[\"month\"] < month]\n",
    "test = data[data[\"month\"] == month]\n",
    "\n",
    "train_data = train.iloc[:, 3:]\n",
    "test_data = test.iloc[:, 3:]\n",
    "\n",
    "train_job_label = train[[\"submit_gpu_job\"]]\n",
    "train_gpu_label = train[[\"submit_gpu_num\"]]\n",
    "test_job_label = test[[\"submit_gpu_job\"]]\n",
    "test_gpu_label = test[[\"submit_gpu_num\"]]\n",
    "\n",
    "train_data.fillna(value=0, inplace=True)\n",
    "test_data.fillna(value=0, inplace=True)\n",
    "\n",
    "# Plot\n",
    "ymax_job = {\"Saturn\": 220, \"Venus\": 50, \"Philly\": 100, \"Uranus\": 100, \"Earth\": 100}\n",
    "ymax_gpu = {\"Saturn\": 1200, \"Venus\": 500, \"Philly\": 150, \"Uranus\": 300, \"Earth\": 200}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpretable Model: Primo(EBM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Job #\"\"\"\n",
    "config = {\"learning_rate\": 0.1, \"interactions\": 10}\n",
    "ebm = PrimoRegressor(model=\"PrAM\", model_config=config, hpo=None)\n",
    "ebm.fit(train_data, train_job_label)\n",
    "pred = ebm.predict(test_data)\n",
    "\n",
    "mae_score = metrics.mean_absolute_error(test_job_label, pred)\n",
    "r2_score = metrics.r2_score(test_job_label, pred)\n",
    "result.at[\"ebm_job_mae\", cluster] = mae_score\n",
    "print(f\"mae_score: {mae_score:.2f}, r2_score: {r2_score:.4f}\")\n",
    "\n",
    "# pred_df.to_csv(f\"./{cluster}_throughput_pred.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Baselines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. LightGBM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Job #\"\"\"\n",
    "model = lgb.LGBMRegressor()\n",
    "model.fit(train_data, train_job_label)\n",
    "pred = model.predict(test_data)\n",
    "\n",
    "mae_score = metrics.mean_absolute_error(test_job_label, pred)\n",
    "r2_score = metrics.r2_score(test_job_label, pred)\n",
    "result.at[\"lgb_job_mae\", cluster] = mae_score\n",
    "print(f\"mae_score: {mae_score:.2f}, r2_score: {r2_score:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. XGBoost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Job #\"\"\"\n",
    "model = xgb.XGBRegressor()\n",
    "model.fit(train_data, train_job_label)\n",
    "pred = model.predict(test_data)\n",
    "\n",
    "mae_score = metrics.mean_absolute_error(test_job_label, pred)\n",
    "r2_score = metrics.r2_score(test_job_label, pred)\n",
    "result.at[\"xgb_job_mae\", cluster] = mae_score\n",
    "print(f\"mae_score: {mae_score:.2f}, r2_score: {r2_score:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Random Forest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Job #\"\"\"\n",
    "model = RandomForestRegressor()\n",
    "model.fit(train_data, train_job_label)\n",
    "pred = model.predict(test_data)\n",
    "\n",
    "mae_score = metrics.mean_absolute_error(test_job_label, pred)\n",
    "r2_score = metrics.r2_score(test_job_label, pred)\n",
    "result.at[\"rf_job_mae\", cluster] = mae_score\n",
    "print(f\"mae_score: {mae_score:.2f}, r2_score: {r2_score:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. DNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Job #\"\"\"\n",
    "model = MLPRegressor()\n",
    "model.fit(train_data, train_job_label)\n",
    "pred = model.predict(test_data)\n",
    "\n",
    "mae_score = metrics.mean_absolute_error(test_job_label, pred)\n",
    "r2_score = metrics.r2_score(test_job_label, pred)\n",
    "result.at[\"dnn_job_mae\", cluster] = mae_score\n",
    "print(f\"mae_score: {mae_score:.2f}, r2_score: {r2_score:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Result Comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prediction Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_predict_all(df, cluster, ymax_job, ymax_gpu, save=False):\n",
    "    fig, (ax1, ax2) = plt.subplots(\n",
    "        ncols=1, nrows=2, constrained_layout=True, figsize=(8, 6)\n",
    "    )\n",
    "    x = np.arange(len(df))\n",
    "    df.reset_index()\n",
    "\n",
    "    if cluster == \"Philly\":\n",
    "        str_month = \"November\"\n",
    "    else:\n",
    "        str_month = \"September\"\n",
    "\n",
    "    ax1.plot(x, df[\"submit_gpu_job\"], linestyle=\"--\", alpha=0.9, label=\"Real\", linewidth=2)\n",
    "    ax1.plot(x, df[\"pred_gpu_job\"], linestyle=\"-\", alpha=0.9, label=\"Prediction\", linewidth=2)\n",
    "    # ax1.set_xlabel(f\"Date in {str_month} ({cluster})\")\n",
    "    ax1.set_ylabel(\"Job Submission\")\n",
    "\n",
    "    tick_interval = 2 * 144  # Two days\n",
    "    ax1.set_xticks(x[::tick_interval])\n",
    "    ax1.set_xticklabels(df.index.day[::tick_interval])\n",
    "    ax1.set_xlim(0, len(df) + 1)\n",
    "    ax1.set_ylim(0, ymax_job[cluster])\n",
    "    ax1.grid(axis=\"y\", linestyle=\":\")\n",
    "    ax1.legend()\n",
    "\n",
    "    ax2.plot(x, df[\"submit_gpu_num\"], linestyle=\"--\", alpha=0.9, label=\"Real\", linewidth=2)\n",
    "    ax2.plot(x, df[\"pred_gpu_num\"], linestyle=\"-\", alpha=0.9, label=\"Prediction\", linewidth=2)\n",
    "    ax2.set_xlabel(f\"Date in {str_month} ({cluster})\")\n",
    "    ax2.set_ylabel(\"GPU Submission\")\n",
    "\n",
    "    ax2.set_xticks(x[::tick_interval])\n",
    "    ax2.set_xticklabels(df.index.day[::tick_interval])\n",
    "    ax2.set_xlim(0, len(df) + 1)\n",
    "    ax2.set_ylim(0, ymax_gpu[cluster])\n",
    "    ax2.grid(axis=\"y\", linestyle=\":\")\n",
    "\n",
    "    sns.despine()\n",
    "\n",
    "    if save:\n",
    "        fig.savefig(f\"./{cluster}_throughput.pdf\", bbox_inches=\"tight\")\n",
    "\n",
    "pred_df = pd.read_csv(f\"./{cluster}_throughput_pred.csv\", parse_dates=[\"time\"], index_col=\"time\")\n",
    "\n",
    "plot_predict_all(pred_df, cluster, ymax_job, ymax_gpu, save=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Interpretation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ebm = ebm.prModel.gam\n",
    "feature_score = ebm.feature_importances_\n",
    "sorted_score = sorted(feature_score, reverse=True)\n",
    "sort_idx = sorted(range(len(feature_score)), key= lambda x: feature_score[x],reverse=True)\n",
    "sorted_feature = [ebm.feature_names[i] for i in sort_idx]\n",
    "\n",
    "sorted_feature = [i.replace(\"_ngpu\", \"\") for i in sorted_feature]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmp = sns.color_palette(\"tab10\")\n",
    "cmp1 = sns.color_palette(\"Blues\",8)\n",
    "def global_interpret(n_features=12, save=False):\n",
    "    fig, ax = plt.subplots(\n",
    "        ncols=1, nrows=1, constrained_layout=True, figsize=(5, 4)\n",
    "    )\n",
    "    \n",
    "    n_features = n_features\n",
    "    x = np.arange(1, n_features+1)\n",
    "\n",
    "    ax.barh(x[::-1], sorted_score[:n_features], label=sorted_feature[:n_features], height=0.2, color=cmp1[5])\n",
    "    ax.scatter(sorted_score[:n_features], x[::-1], s=50, color=cmp1[5])\n",
    "    \n",
    "    ax.set_xlim(0, 1)\n",
    "    ax.set_yticks(x)\n",
    "    ax.set_yticklabels(sorted_feature[:n_features][::-1])\n",
    "    ax.set_xlabel(f\"Average Absolute Score\")\n",
    "    ax.grid(axis=\"x\", linestyle=\":\")\n",
    "    sns.despine()\n",
    "\n",
    "    if save:\n",
    "        fig.savefig(f\"./interpret_{cluster}_throughput.pdf\", bbox_inches=\"tight\")\n",
    "\n",
    "\n",
    "global_interpret(n_features=12, save=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.isotonic import IsotonicRegression\n",
    "\n",
    "feature_index = 2\n",
    "x = np.array(range(len(ebm.additive_terms_[feature_index])))\n",
    "y = ebm.additive_terms_[feature_index]\n",
    "w = ebm.preprocessor_.col_bin_counts_[feature_index]\n",
    "direction = \"auto\"\n",
    "ir = IsotonicRegression(out_of_bounds=\"clip\", increasing=direction)\n",
    "y_ = ir.fit_transform(x, y, sample_weight=w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hour_interpret(save=False):\n",
    "    fig, ax = plt.subplots(\n",
    "        ncols=1, nrows=1, constrained_layout=True, figsize=(5, 4)\n",
    "    )\n",
    "    y2 = 0\n",
    "    ax.plot(x, y, '-', label=\"Shape Function\", marker=\"X\", markersize=8, linewidth=2.5, color=cmp1[5], alpha=1)\n",
    "\n",
    "    ax.set_xlabel(\"Hour Bins\")\n",
    "    ax.set_ylabel(f\"Score\")\n",
    "    ax.set_xlim(0, 24)\n",
    "    ax.set_xticks([0, 4, 8, 12, 16, 20, 24])\n",
    "    # ax.set_ylim(-1.5, 4)\n",
    "    sns.despine()\n",
    "\n",
    "    ax.grid(axis=\"y\", linestyle=\":\")\n",
    "    ax.axhline(y=0, c=cmp1[-1], linewidth=2)\n",
    "    ax.fill_between(x, y, y2, facecolor=cmp1[5], alpha = 0.1, hatch='/')\n",
    "    ax.legend(loc=0, fontsize=16)\n",
    "\n",
    "    if save:\n",
    "        fig.savefig(f\"./interpret_{cluster}_shapefunc.pdf\", bbox_inches=\"tight\")\n",
    "\n",
    "\n",
    "hour_interpret(save=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "784914556ea7aafefbdcb0c4fefea2600a01efb6a6b7916d4154dc17a3e6434f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
